package com.moon.java.code.analyzer.test;

import java.util.ArrayList;
import java.util.List;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.NotFoundException;
import javassist.bytecode.BadBytecode;
import javassist.bytecode.BootstrapMethodsAttribute;
import javassist.bytecode.BootstrapMethodsAttribute.BootstrapMethod;
import javassist.bytecode.CodeAttribute;
import javassist.bytecode.CodeIterator;
import javassist.bytecode.ConstPool;
import javassist.bytecode.MethodInfo;
import javassist.bytecode.Opcode;

public class LambdaMethodAnalyzer2 {
    public static void main(String[] args) {
        try {
            // 获取类池
            ClassPool pool = ClassPool.getDefault();
            // 获取要分析的类
            CtClass cc = pool.get("com.moon.java.code.analyzer.test.LambdaTest");

            // 查找所有以 lambda$ 开头的方法
            List<CtMethod> lambdaMethods = new ArrayList<>();
            CtMethod[] allMethods = cc.getDeclaredMethods();
            for (CtMethod m : allMethods) {
                if (m.getName().startsWith("lambda$")) {
                    lambdaMethods.add(m);
                    System.out.println("找到Lambda方法: " + m.getName());
                }
            }

            // 获取 testMethod 方法
            CtMethod testMethod = cc.getDeclaredMethod("testMethod");
            MethodInfo methodInfo = testMethod.getMethodInfo();

            // 从类文件获取BootstrapMethods属性
            BootstrapMethodsAttribute bsmAttr = (BootstrapMethodsAttribute) cc.getClassFile().getAttribute(BootstrapMethodsAttribute.tag);

            if (bsmAttr == null) {
                System.out.println("警告：类文件中没有BootstrapMethods属性");
                return;
            }

            CodeAttribute codeAttribute = methodInfo.getCodeAttribute();
            if (codeAttribute == null) {
                System.out.println("警告：testMethod没有代码属性");
                return;
            }

            ConstPool constPool = methodInfo.getConstPool();
            CodeIterator iterator = codeAttribute.iterator();

            while (iterator.hasNext()) {
                int index = iterator.next();
                int opcode = iterator.byteAt(index);

                if (opcode == Opcode.INVOKEDYNAMIC) {
                    // 读取 invokedynamic 的操作数（指向常量池中的 InvokeDynamic 条目）
                    int invokeDynamicIndex = iterator.u16bitAt(index + 1);

                    // 从常量池获取真正的 bootstrap 方法索引
                    int bsmIndex = constPool.getInvokeDynamicBootstrap(invokeDynamicIndex);
                    System.out.println("发现invokedynamic指令，常量池索引=" + invokeDynamicIndex + ", bsmIndex=" + bsmIndex);

                    // 检查索引是否有效
                    if (bsmIndex >= bsmAttr.getMethods().length) {
                        System.out.println("错误：bootstrap方法索引越界，最大应为 " + (bsmAttr.getMethods().length - 1));
                        continue;
                    }

                    BootstrapMethod bsmMethod = bsmAttr.getMethods()[bsmIndex];
                    CtMethod lambdaMethod = findLambdaMethod(cc, bsmMethod, constPool, lambdaMethods);

                    if (lambdaMethod != null) {
                        System.out.println("方法 testMethod 调用了 Lambda 方法: " + lambdaMethod.getName());
                        analyzeCalledMethodsInLambda(lambdaMethod);
                    } else {
                        System.out.println("未找到对应的Lambda方法");
                    }
                }
            }

        } catch (NotFoundException | BadBytecode e) {
            e.printStackTrace();
        }
    }

    private static CtMethod findLambdaMethod(CtClass cc, BootstrapMethod bsmMethod, ConstPool constPool, List<CtMethod> lambdaMethods) throws NotFoundException {
        // 获取Bootstrap方法的参数
        int[] args = bsmMethod.arguments;

        for (int arg : args) {
            if (constPool.getTag(arg) == ConstPool.CONST_MethodHandle) {
                int refKind = constPool.getMethodHandleKind(arg);
                int refIndex = constPool.getMethodHandleIndex(arg);

                System.out.println("处理方法句柄: kind=" + refKind + ", index=" + refIndex);

                // REF_invokeStatic=6, REF_invokeSpecial=5
                if (refKind == ConstPool.REF_invokeSpecial || refKind == ConstPool.REF_invokeStatic) {
                    String methodName = constPool.getMethodrefName(refIndex);
                    String className = constPool.getMethodrefClassName(refIndex);

                    System.out.println("查找方法: " + className + "." + methodName);

                    // 在lambda方法列表中查找匹配的方法
                    for (CtMethod lambdaMethod : lambdaMethods) {
                        if (lambdaMethod.getName().equals(methodName)) {
                            return lambdaMethod;
                        }
                    }
                }
            }
        }
        return null;
    }

    private static void analyzeCalledMethodsInLambda(CtMethod lambdaMethod) throws BadBytecode, NotFoundException {
        MethodInfo methodInfo = lambdaMethod.getMethodInfo();
        CodeAttribute codeAttribute = methodInfo.getCodeAttribute();

        if (codeAttribute == null) {
            System.out.println("警告：Lambda方法没有代码属性");
            return;
        }

        ConstPool constPool = methodInfo.getConstPool();
        CodeIterator iterator = codeAttribute.iterator();
        List<String> calledMethods = new ArrayList<>();

        while (iterator.hasNext()) {
            int index = iterator.next();
            int opcode = iterator.byteAt(index);

            if (opcode == Opcode.INVOKEVIRTUAL || opcode == Opcode.INVOKESPECIAL || opcode == Opcode.INVOKESTATIC || opcode == Opcode.INVOKEINTERFACE) {

                int methodrefIndex = iterator.u16bitAt(index + 1);
                String className = constPool.getClassInfo(constPool.getMemberClass(methodrefIndex));
                String methodName = constPool.getMethodrefName(methodrefIndex);
                calledMethods.add(className + "." + methodName);
            }
        }

        System.out.println("Lambda方法 " + lambdaMethod.getName() + " 中调用的方法:");
        for (String methodName : calledMethods) {
            System.out.println("  " + methodName);
        }
    }
}